# ElGamal.py
# ElGamal加密算法
# -*- coding: utf-8 -*-
"""
Created on 2021

@author: Ximing
"""

import math
import random


def elgaml_encrypt(y, g, p, M):
    ''''公钥加密: 加密信息M,返回密文对C1,C2'''
    pt = message
    C = [[] for i in range(2)]
    for i in pt:
        r = random.randint(0, p - 1)
        while 1 != gcd(r, p - 1):
            r = random.randint(0, p - 1)
        C[0].append((g ** r) % p)
        C[1].append(((y ** r) * ord(i) % p))
    return C

def elgaml_decrypt(x, p, C):
    '''''私钥解密: 解密密文对C1, C2得到信息M'''
    plaintext = []
    temp = 0
    for i in range(0, len(C[0])):
        C1_temp = (C[0][i] ** x) % p
        C1_inverse = extend_gcd(C1_temp, p)[0]
        M = (C[1][i] * C1_inverse) % p
        plaintext.append(chr(M))
    return "".join(plaintext)


def extend_gcd(dividend, divisor):
    '''''求乘法逆元:返回各自逆元和最大公约数'''
    if divisor == 0:
        return 1, 0, dividend
    x2, y2, remainder = extend_gcd(divisor, dividend % divisor)
    temp = x2
    x1 = y2
    y1 = temp - (dividend // divisor) * y2
    return x1, y1, remainder


def getPrime():
    '''''定义获取素数的函数'''
    ptmp = random.randint(260, 400)
    while not is_prime(ptmp):
        ptmp = random.randint(260, 400)
    return ptmp


def get_public_key(x, g, p):
    '''''得到公钥y'''
    y = (g ** x) % p
    return y


def is_prime(num):
    '''''判断是否为素数'''
    if num == 2 or num == 3:
        return True
    if num % 6 != 1 and num % 6 != 5:
        return False
    for i in range(5, int(math.sqrt(num)) + 1, 6):
        if num % i == 0 or num % (i + 2) == 0:
            return False
    return True


def gcd(m, n):
    '''求最大公约数'''
    if m < n:
        m, n = n, m
    if n == 1 or m == n + 1 or m == 2 * n + 1 or m == 2 * n - 1 or is_prime(m) == True:
        return 1
    if n == 0:
        return m
    else:
        return gcd(n, m % n)


def get_prime_factor(num):
    '''''得到不同的素因子'''
    prime_factor = []
    while num % 2 == 0:
        if 2 not in prime_factor:
            prime_factor.append(2)
        num //= 2
    for i in range(3, int(math.sqrt(num)) + 1, 2):
        while num % i == 0:
            if i not in prime_factor:
                prime_factor.append(i)
            num //= i
    if num > 2:
        prime_factor.append(num)
    return prime_factor


def get_primitive_root(num):
    '''''得到原根的最小值'''
    prime_list = get_prime_factor(num - 1)
    for g in range(2, num):
        flag = 0
        for i in prime_list:
            if g ** ((num - 1) // i) % num == 1:
                flag = 1
                break
        if flag == 0:
            return g
    return False


def init():
    p = getPrime()
    g = get_primitive_root(p)
    x = random.randint(20, p - 2)
    return x, g, p


if __name__ == "__main__":
    x, g, p = init()
    y = get_public_key(x, g, p)
    print('公钥:(%d,原根 = %d,素数 = %d)' % (y, g, p))
    print('私钥:(%d)' % (x))
    message = input('传送的信息M:')
    C = elgaml_encrypt(y, g, p, message)
    print('公钥加密:加密信息M,返回密文对C1,C2:', C)
    d_message = elgaml_decrypt(x, p, C)
    print('私钥解密: 解密密文对C1, C2得到信息M:', d_message)
